#!/usr/bin/env python3
"""
multi_extract.py
多格式解压工具（支持 zip, tar, tar.gz, tar.bz2, gz（单文件）, rar, 7z, ace via patool if installed）

特点：
- 自动识别格式并选择合适的解压方式
- 防止路径穿越（zip slip）
- 支持提取到指定目录
- 支持递归解压嵌套压缩包（可选）
- 命令行界面（支持 --verbose --recursive）

依赖：
- 内置：shutil, tarfile, zipfile, gzip
- 可选：py7zr (7z 支持), rarfile (rar 支持), patool (额外格式)
  安装：pip install py7zr rarfile patool
  注意：rarfile 可能还需要系统级的 unrar / unar 可执行文件。

示例：
  python3 multi_extract.py archive.zip -d out_dir --recursive

"""

import argparse
import os
import shutil
import sys
import tarfile
import zipfile
import pyzipper
import gzip
from pathlib import Path
import tempfile

# 可选依赖尝试导入
try:
    import py7zr
except Exception as e:
    print(e)
    py7zr = None

try:
    import rarfile
except Exception as e:
    print(e)
    rarfile = None

try:
    import patoolib
except Exception as e:
    print(e)
    patoolib = None

def extract_filename_from_path(file_path):
    """
    从文件路径中提取文件名
    
    Args:
        file_path (str): 完整的文件路径
        
    Returns:
        str: 文件名（包含扩展名）
    """
    try:
        # 方法1: 使用os.path.basename
        filename = os.path.basename(file_path)
        return filename
    except Exception as e:
        print(f"提取文件名时出错: {e}")
        return None

def get_all_files(directory, exclude_patterns = ['ht.json']):
    path = Path(directory)
    files = [str(f) for f in path.glob('**/*') 
             if f.is_file() and f.name not in exclude_patterns]
    return files


def safe_extract_tar(tar: tarfile.TarFile, path: str):
    for member in tar.getmembers():
        member_path = os.path.join(path, member.name)
        abs_path = os.path.abspath(member_path)
        if not abs_path.startswith(os.path.abspath(path) + os.sep) and abs_path != os.path.abspath(path):
            raise Exception(f"尝试解压到非法路径: {member.name}")
    tar.extractall(path)


def safe_extract_zipper(zf: zipfile.ZipFile, path: str):
    with pyzipper.AESZipFile(zf) as z:
        zf.extractall(path)

def safe_extract_zip(zf: zipfile.ZipFile, path: str):
    for member in zf.namelist():
        # try:
        #     member = member.encode("cp437").decode("gbk")
        # except UnicodeDecodeError:
        #     pass
        member_path = os.path.join(path, member)
        abs_path = os.path.abspath(member_path)
        print(f"member: {member}, {member_path} , {abs_path}")
        if not abs_path.startswith(os.path.abspath(path) + os.sep) and abs_path != os.path.abspath(path):
            raise Exception(f"尝试解压到非法路径: {member}")
    zf.extractall(path)


def extract_gz_file(gz_path: str, dest_dir: str):
    # .gz 很多时候仅是单文件压缩（例如 .tar.gz 外层）
    out_name = os.path.splitext(os.path.basename(gz_path))[0]
    out_path = os.path.join(dest_dir, out_name)
    with gzip.open(gz_path, 'rb') as f_in, open(out_path, 'wb') as f_out:
        shutil.copyfileobj(f_in, f_out)
    return out_path


def is_archive(filepath: str) -> bool:
    # 简单判断：基于扩展名
    lower = filepath.lower()
    exts = ['.zip', '.tar', '.gz', '.tgz', '.bz2', '.tbz', '.xz', '.7z', '.rar', '.tar.gz', '.tar.bz2']
    return any(lower.endswith(e) for e in exts)


def extract_archive(filepath: str, dest_dir: str = None, recursive: bool = False, verbose: bool = False):
    
    if dest_dir is None:
        dir = os.path.splitext(filepath)[:-1]
        dest_dir = os.path.join(os.getcwd(), dir[0] + '_extracted')
        print(f"dest_dir: {dest_dir}")
        Path(dest_dir).mkdir(parents=True, exist_ok=True)
    lower = filepath.lower()

    try:
        if lower.endswith('.zip'):
            if verbose: print(f"Unpacking ZIP: {filepath} -> {dest_dir}")
            with zipfile.ZipFile(filepath, 'r') as zf:
                safe_extract_zip(zf, dest_dir)

        elif lower.endswith('.tar') or lower.endswith('.tgz') or lower.endswith('.tar.gz') or lower.endswith('.tar.bz2') or lower.endswith('.tbz') or lower.endswith('.tar.xz'):
            if verbose: print(f"Unpacking TAR: {filepath} -> {dest_dir}")
            mode = 'r'
            if lower.endswith('.tgz') or lower.endswith('.tar.gz'):
                mode = 'r:gz'
            elif lower.endswith('.tar.bz2') or lower.endswith('.tbz'):
                mode = 'r:bz2'
            elif lower.endswith('.tar.xz'):
                mode = 'r:xz'
            with tarfile.open(filepath, mode) as tf:
                safe_extract_tar(tf, dest_dir)

        elif lower.endswith('.gz') and not lower.endswith('.tar.gz'):
            if verbose: print(f"Unpacking GZ(single file): {filepath} -> {dest_dir}")
            extract_gz_file(filepath, dest_dir)

        elif lower.endswith('.7z'):
            if py7zr is None:
                raise RuntimeError('缺少 py7zr 库，无法解压 7z 文件。请运行: pip install py7zr')
            if verbose: print(f"Unpacking 7z: {filepath} -> {dest_dir}")
            with py7zr.SevenZipFile(filepath, mode='r') as archive:
                archive.extractall(path=dest_dir)

        elif lower.endswith('.rar'):
            if rarfile is None:
                raise RuntimeError('缺少 rarfile 库，无法解压 rar 文件。请运行: pip install rarfile，并确保系统安装了 unrar/unar')
            if verbose: print(f"Unpacking RAR: {filepath} -> {dest_dir}")
            try:
                with rarfile.RarFile(filepath) as rf:
                    rf.extractall(dest_dir)
            except Exception:
                print("rarfile 解压失败使用 patoolib")
                patoolib.extract_archive(filepath, outdir=dest_dir)

        else:
            # 退回到 patool（支持更多格式），如果安装了
            if patoolib is not None:
                if verbose: print(f"Using patool to extract: {filepath} -> {dest_dir}")
                patoolib.extract_archive(filepath, outdir=dest_dir)
            else:
                raise RuntimeError(f"不支持的压缩格式或缺少依赖: {filepath}")

        # 如果递归模式，扫描目标目录，继续解压子归档
        if recursive:
            for root, dirs, files in os.walk(dest_dir):
                for name in files:
                    full = os.path.join(root, name)
                    if is_archive(full) and full != filepath:
                        subdir = full + '.extracted'
                        if verbose: print(f"递归解压: {full} -> {subdir}")
                        extract_archive(full, subdir, recursive=recursive, verbose=verbose)
        return dest_dir
    except Exception as e:
        raise


def main():
    parser = argparse.ArgumentParser(description='多格式解压工具 (zip, tar, gz, 7z, rar...)')
    parser.add_argument('archive', help='要解压的文件路径')
    parser.add_argument('-d', '--dest', help='输出目录，默认是当前目录下的 archive_name_extracted', default=None)
    parser.add_argument('--recursive', action='store_true', help='递归解压嵌套的压缩包')
    parser.add_argument('--verbose', action='store_true', help='输出详细日志')

    args = parser.parse_args()
    archive = args.archive

    if not os.path.isfile(archive):
        print('找不到文件:', archive)
        sys.exit(2)

    if args.dest:
        dest = args.dest
    else:
        base = os.path.splitext(os.path.basename(archive))[0]
        dest = os.path.join(os.getcwd(), base + '_extracted')
    print(f"{archive, dest, args.recursive, args.verbose}")
    try:
        extract_archive(archive, None, recursive=args.recursive, verbose=args.verbose)
        print('解压完成 ->', dest)
    except Exception as e:
        print('解压失败:', str(e))
        sys.exit(1)


if __name__ == '__main__':
    main()
